from torchvision import datasets, transforms
from torch.utils.data import DataLoader

def get_cifar10_data_loaders(batch_size=64):
    """
    获取CIFAR10数据集的数据加载器，包含增强的数据预处理
    
    参数:
        batch_size: 每个批次的样本数量
        
    返回:
        train_loader: 训练集数据加载器
        test_loader: 测试集数据加载器
    """
    # 训练集数据预处理和增强
    train_transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),  # 随机水平翻转
        transforms.RandomRotation(15),      # 随机旋转±15度
        transforms.RandomAffine(0, translate=(0.1, 0.1)),  # 随机平移
        transforms.RandomCrop(32, padding=4),  # 随机裁剪
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),  # 颜色抖动
        transforms.ToTensor(),  # 转换为张量
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))  # CIFAR10专用归一化
    ])
    
    # 测试集数据预处理(不进行增强)
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
    ])

    # 加载训练集和测试集
    train_dataset = datasets.CIFAR10(
        root='./data', 
        train=True, 
        download=True, 
        transform=train_transform
    )
    test_dataset = datasets.CIFAR10(
        root='./data', 
        train=False, 
        download=True, 
        transform=test_transform
    )

    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset, 
        batch_size=batch_size, 
        shuffle=True,
        num_workers=2,  # 使用多线程加载数据
        pin_memory=True  # 加速GPU传输
    )
    test_loader = DataLoader(
        test_dataset, 
        batch_size=batch_size, 
        shuffle=False,
        num_workers=2,
        pin_memory=True
    )

    return train_loader, test_loader